"""
PyG完全兼容DGL的图结构实现

这个模块提供了完全兼容DGL的图结构包装器，封装PyG的Data/Batch对象，
确保所有API和返回值形状与DGL一致。
"""

import torch
import torch_npu
from torch_npu.contrib import transfer_to_npu
from typing import Dict, Tuple, List, Optional, Union, Any, TypeVar, Protocol
from torch import Tensor

# 条件导入PyG
try:
    from torch_geometric.data import Data as PyGData
    from torch_geometric.data import Batch as PyGBatch
    from torch_scatter import scatter_mean, scatter_max, scatter_sum
    HAS_PYG = True
except ImportError:
    HAS_PYG = False
    PyGData = Any
    PyGBatch = Any

# 条件导入DGL
try:
    import dgl
    from dgl import DGLGraph
    from dgl import batch as dgl_batch
    HAS_DGL = True
except ImportError:
    HAS_DGL = False
    DGLGraph = Any
    dgl_batch = None

# 本地模块导入
from .graph import is_dgl_graph

# 类型变量
T = TypeVar('T', bound=Tensor)


class NodeDataView:
    """完全兼容DGL的NodeDataView实现"""

    def __init__(self, graph):
        """初始化节点数据视图
        
        参数:
            graph: CompatGraph对象
        """
        self._graph = graph
        self._pyg_data = graph._pyg_data

    def __getitem__(self, key: str) -> Tensor:
        """获取节点特征，确保维度和格式与DGL完全一致
        
        参数:
            key: 特征名称
            
        返回:
            节点特征张量
        """
        if key in self._pyg_data:
            return self._pyg_data[key]
        raise KeyError(f"节点特征 '{key}' 不存在")

    def __setitem__(self, key: str, value: Tensor):
        """设置节点特征，确保维度和格式与DGL完全一致
        
        参数:
            key: 特征名称
            value: 特征值张量
        """
        self._pyg_data[key] = value

    def keys(self):
        """返回所有节点特征的键
        
        返回:
            节点特征键列表
        """
        return [k for k in self._pyg_data.keys() if not k.startswith('edge_') and k != 'edge_index']

    def __contains__(self, key):
        """检查节点特征是否存在
        
        参数:
            key: 特征名称
            
        返回:
            布尔值，表示特征是否存在
        """
        return key in self._pyg_data


class EdgeDataView:
    """完全兼容DGL的EdgeDataView实现"""

    def __init__(self, graph):
        """初始化边数据视图
        
        参数:
            graph: CompatGraph对象
        """
        self._graph = graph
        self._pyg_data = graph._pyg_data

    def __getitem__(self, key: str) -> Tensor:
        """获取边特征，确保维度和格式与DGL完全一致
        
        参数:
            key: 特征名称
            
        返回:
            边特征张量
        """
        edge_key = f"edge_{key}"
        if edge_key in self._pyg_data:
            return self._pyg_data[edge_key]
        raise KeyError(f"边特征 '{key}' 不存在")

    def __setitem__(self, key: str, value: Tensor):
        """设置边特征，确保维度和格式与DGL完全一致
        
        参数:
            key: 特征名称
            value: 特征值张量
        """
        self._pyg_data[f"edge_{key}"] = value

    def keys(self):
        """返回所有边特征的键
        
        返回:
            边特征键列表
        """
        return [k[5:] for k in self._pyg_data.keys() if k.startswith('edge_') and k != 'edge_index']

    def __contains__(self, key):
        """检查边特征是否存在
        
        参数:
            key: 特征名称
            
        返回:
            布尔值，表示特征是否存在
        """
        return f"edge_{key}" in self._pyg_data


class CompatGraph:
    """
    完全兼容DGL的图结构包装器，封装PyG的Data/Batch对象
    
    提供与DGL完全一致的API，同时确保所有操作都保持正确的张量形状
    """
    
    def __init__(self, pyg_data):
        """初始化兼容图对象
        
        参数:
            pyg_data: PyG的Data或Batch对象
        """
        if not HAS_PYG:
            raise ImportError("未安装PyG，无法创建CompatGraph对象")
            
        self._pyg_data = pyg_data
        self._ndata = NodeDataView(self)
        self._edata = EdgeDataView(self)
        self._cache = {}  # 缓存常用计算结果
        
    @property
    def ndata(self):
        """节点数据访问接口，与DGL的graph.ndata完全一致"""
        return self._ndata
        
    @property
    def edata(self):
        """边数据访问接口，与DGL的graph.edata完全一致"""
        return self._edata
    
    def number_of_nodes(self) -> int:
        """返回图中节点数量
        
        返回:
            节点数量
        """
        return self._pyg_data.num_nodes
        
    def number_of_edges(self) -> int:
        """返回图中边数量
        
        返回:
            边数量
        """
        return self._pyg_data.num_edges
    
    def batch_num_nodes(self) -> Tensor:
        """返回批处理中每个图的节点数量
        
        返回:
            每个图的节点数量张量
        """
        # 尝试从缓存获取
        cache_key = 'batch_num_nodes'
        if cache_key in self._cache:
            return self._cache[cache_key]
            
        # 计算并缓存结果
        if hasattr(self._pyg_data, 'batch') and self._pyg_data.batch is not None:
            unique, counts = torch.unique(self._pyg_data.batch, return_counts=True)
            self._cache[cache_key] = counts
            return counts
        else:
            result = torch.tensor([self._pyg_data.num_nodes], device=self.device)
            self._cache[cache_key] = result
            return result
    
    def batch_num_edges(self) -> Tensor:
        """返回批处理中每个图的边数量
        
        返回:
            每个图的边数量张量
        """
        # 尝试从缓存获取
        cache_key = 'batch_num_edges'
        if cache_key in self._cache:
            return self._cache[cache_key]
            
        # PyG没有直接存储批次边信息，需要基于节点批次推导
        if hasattr(self._pyg_data, 'batch') and self._pyg_data.batch is not None:
            src, dst = self._pyg_data.edge_index
            src_batch = self._pyg_data.batch[src]
            unique, counts = torch.unique(src_batch, return_counts=True)
            self._cache[cache_key] = counts
            return counts
        else:
            result = torch.tensor([self._pyg_data.num_edges], device=self.device)
            self._cache[cache_key] = result
            return result
    
    def edges(self) -> Tuple[Tensor, Tensor]:
        """返回图中所有边的源节点和目标节点索引
        
        返回:
            (源节点索引张量, 目标节点索引张量)
        """
        # 尝试从缓存获取
        cache_key = 'edges'
        if cache_key in self._cache:
            return self._cache[cache_key]
            
        # 计算并缓存结果
        src, dst = self._pyg_data.edge_index
        result = (src, dst)
        self._cache[cache_key] = result
        return result
    
    def has_node(self, node_id: int) -> bool:
        """检查指定节点是否存在
        
        参数:
            node_id: 节点ID
            
        返回:
            布尔值，表示节点是否存在
        """
        return 0 <= node_id < self.number_of_nodes()
    
    def has_edges_between(self, u, v) -> bool:
        """检查两个节点之间是否有边
        
        参数:
            u: 源节点ID
            v: 目标节点ID
            
        返回:
            布尔值，表示是否存在边
        """
        src, dst = self.edges()
        mask = (src == u) & (dst == v)
        return torch.any(mask).item()
    
    def in_degrees(self, v: Optional[int] = None) -> Union[int, Tensor]:
        """返回节点的入度
        
        参数:
            v: 可选的节点ID，如果提供则返回该节点的入度
            
        返回:
            如果v不为None，返回该节点的入度；否则返回所有节点的入度张量
        """
        _, dst = self.edges()
        if v is not None:
            return torch.sum(dst == v).item()
        else:
            return torch.bincount(dst, minlength=self.number_of_nodes())
    
    def out_degrees(self, u: Optional[int] = None) -> Union[int, Tensor]:
        """返回节点的出度
        
        参数:
            u: 可选的节点ID，如果提供则返回该节点的出度
            
        返回:
            如果u不为None，返回该节点的出度；否则返回所有节点的出度张量
        """
        src, _ = self.edges()
        if u is not None:
            return torch.sum(src == u).item()
        else:
            return torch.bincount(src, minlength=self.number_of_nodes())
    
    def to(self, device):
        """将图移动到指定设备
        
        参数:
            device: 目标设备
            
        返回:
            移动到目标设备的图对象
        """
        # 清除缓存
        self._cache = {}
        
        # 移动PyG数据到目标设备
        self._pyg_data = self._pyg_data.to(device)
        
        return self
    
    @property
    def device(self):
        """返回图所在设备
        
        返回:
            图所在设备
        """
        if hasattr(self._pyg_data, 'edge_index'):
            return self._pyg_data.edge_index.device
        return torch.device('cpu')
    
    @property
    def batch_size(self) -> int:
        """返回批处理大小
        
        返回:
            批处理大小
        """
        if hasattr(self._pyg_data, 'batch') and self._pyg_data.batch is not None:
            return int(self._pyg_data.batch.max().item() + 1)
        return 1
    
    def __repr__(self):
        """返回图的字符串表示
        
        返回:
            图的字符串表示
        """
        return (f"CompatGraph(num_nodes={self.number_of_nodes()}, "
                f"num_edges={self.number_of_edges()}, "
                f"batch_size={self.batch_size})")


def compat_batch(graph_list: List[CompatGraph]) -> CompatGraph:
    """
    将图列表合并为一个批处理图，与dgl.batch行为完全一致
    
    参数:
        graph_list: CompatGraph对象列表
        
    返回:
        批处理后的CompatGraph对象
    """
    if not HAS_PYG:
        raise ImportError("未安装PyG，无法执行批处理操作")
        
    if not graph_list:
        raise ValueError("图列表为空")
        
    # 获取底层PyG数据对象
    pyg_data_list = [g._pyg_data for g in graph_list]
    
    # 使用PyG的Batch.from_data_list批处理图
    batched_data = PyGBatch.from_data_list(pyg_data_list)
    
    # 创建批处理图
    batched_graph = CompatGraph(batched_data)
    
    return batched_graph


def compat_unbatch(batched_graph: CompatGraph) -> List[CompatGraph]:
    """
    将批处理图拆分为图列表，与dgl.unbatch行为一致
    
    参数:
        batched_graph: CompatGraph批处理对象
        
    返回:
        CompatGraph对象列表
    """
    if not HAS_PYG:
        raise ImportError("未安装PyG，无法执行解批处理操作")
        
    if not hasattr(batched_graph._pyg_data, 'batch') or batched_graph._pyg_data.batch is None:
        # 单图情况，直接返回包含自身的列表
        return [batched_graph]
        
    # 使用PyG的Batch.to_data_list解批处理
    from torch_geometric.data import Batch
    pyg_graphs = Batch.to_data_list(batched_graph._pyg_data)
    
    # 转换为CompatGraph对象
    return [CompatGraph(g) for g in pyg_graphs]


def is_compat_graph(graph) -> bool:
    """
    检查对象是否为CompatGraph
    
    参数:
        graph: 要检查的对象
        
    返回:
        布尔值，表示对象是否为CompatGraph
    """
    return isinstance(graph, CompatGraph)


def convert_to_compat_graph(graph) -> CompatGraph:
    """
    将图对象转换为CompatGraph
    
    参数:
        graph: DGL图、PyG图或其他图对象
        
    返回:
        CompatGraph对象
    """
    if is_compat_graph(graph):
        return graph
        
    if is_dgl_graph(graph):
        # DGL图转换为PyG图
        if not HAS_PYG:
            raise ImportError("未安装PyG，无法进行转换")
            
        data = PyGData()
        src, dst = graph.edges()
        data.edge_index = torch.stack([src, dst])
        
        # 复制节点特征
        for key in graph.ndata:
            data[key] = graph.ndata[key]
        
        # 复制边特征
        for key in graph.edata:
            data[f"edge_{key}"] = graph.edata[key]
            
        return CompatGraph(data)
    elif isinstance(graph, PyGData) or isinstance(graph, PyGBatch):
        # PyG图直接包装
        return CompatGraph(graph)
    else:
        # 尝试获取边索引
        if hasattr(graph, 'edge_index'):
            return CompatGraph(graph)
        
        raise TypeError(f"不支持的图类型: {type(graph)}") 